package com.chb.flink.sink

import java.text.SimpleDateFormat
import java.time.LocalDate
import java.time.format.DateTimeFormatter
import java.util.Date

import com.chb.flink.source.{MyCustomerSource, StationLog}
import org.apache.flink.api.common.serialization.SimpleStringEncoder
import org.apache.flink.core.fs.Path
import org.apache.flink.streaming.api.functions.sink.filesystem.StreamingFileSink
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.{DefaultRollingPolicy, OnCheckpointRollingPolicy}
import org.apache.flink.streaming.api.scala.StreamExecutionEnvironment

object HDFSFileSink {
    def main(args: Array[String]): Unit = {
        //初始化Flink的Streaming（流计算）上下文执行环境
        val streamEnv = StreamExecutionEnvironment.getExecutionEnvironment
        streamEnv.setParallelism(1)
        //导入隐式转换
        import org.apache.flink.streaming.api.scala._

        streamEnv.enableCheckpointing(1000)

        // 数据源
        val data: DataStream[StationLog] = streamEnv.addSource(new MyCustomerSource)

        // 创建一个文件滚动策略
        val rollingPolicy: DefaultRollingPolicy[StationLog, String] = DefaultRollingPolicy.create()
            .withRolloverInterval(2000) // 两秒生产一个文件
            .withInactivityInterval(1000) // 不活动的间隔
            .withMaxPartSize(5 * 1024 * 1024)
            .build()

        // 创建sink
        val hdfsSink = StreamingFileSink.forRowFormat[StationLog](
            new Path("hdfs://hadoop01:9000/sink001/"), new SimpleStringEncoder[StationLog]("UTF-8")
        )
            .withBucketCheckInterval(1000) // 检查分桶间隔
            .withBucketAssigner(new MemberBucketAssigner)
            .withRollingPolicy(rollingPolicy)
            .build()

        // 添加sink
        data.addSink(hdfsSink)

        streamEnv.execute()
    }


    import org.apache.flink.core.io.SimpleVersionedSerializer
    import org.apache.flink.streaming.api.functions.sink.filesystem.BucketAssigner
    import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.SimpleVersionedStringSerializer

    /**
     * 自定义分桶策略
     */
    class MemberBucketAssigner extends BucketAssigner[StationLog, String] {
        // 指定桶名 yyyy-mm-dd/HH
        override def getBucketId(info: StationLog, context: BucketAssigner.Context): String = {
            val date = new Date(info.callTime)
            new SimpleDateFormat("yyyy-MM-dd/HH").format(date)
        }

        override def getSerializer: SimpleVersionedSerializer[String] = SimpleVersionedStringSerializer.INSTANCE
    }

}

